package com.cwj.algorithm.tree;

import java.util.LinkedList;
import java.util.Queue;

/**
 * @author chenwujie
 * @date 2020/11/17 8:41
 */
public class BinarySearchTree<K extends Comparable<? super K>, V> {
    private BinaryNode<K, V> root;
    private int count;

    public BinarySearchTree(){
    }

    /**
     * 添加节点
     */
    public void put(K key, V value){
        if(root == null) {
            this.root = new BinaryNode<>(key, value, null, null);
            count++;
        }else{
            put(root, key, value);
        }
    }

    private void put(BinaryNode<K, V> root, K key, V value){
        BinaryNode<K, V> newNode = new BinaryNode<>(key, value, null, null);
        int i = key.compareTo(root.key);
        if(i == 0){
            root.value = value;
        }else if(i < 0){
            if(root.left == null){
                root.left = newNode;
                count++;
            }else{
                put(root.left, key, value);
            }
        }else{
            if(root.right == null){
                root.right = newNode;
                count++;
            }else{
                put(root.right, key, value);
            }
        }
    }

    /**
     * 获取节点值
     */
    public V get(K key){
        return get(root, key);
    }
    private V get(BinaryNode<K, V> x, K key){
        if(x == null){
            return null;
        }

        int i = key.compareTo(x.key);
        if(i == 0){
            return x.value;
        }else if(i < 0){
            return get(x.left, key);
        }else{
            return get(x.right, key);
        }
    }

    public void remove(K key){
        remove(root, key);
    }
    private BinaryNode<K, V> remove(BinaryNode<K, V> x, K key){
        if(x == null){
            return null;
        }
        int i = key.compareTo(x.key);
        if(i == 0){
            count--;
            if(x.right == null){
                return x.left;
            }
            if(x.left == null){
                return x.right;
            }
            // 获取右节点最小节点
            BinaryNode<K, V> minNode = x.right;
            while(minNode.left != null){
                minNode = minNode.left;
            }
            // 断开右节点最小节点
            BinaryNode<K, V> n = x.right;
            while(n.left != null){
                if(n.left.left == null){
                    n.left = null;
                }else{
                    n = n.left;
                }
            }
            // 最小右节点的子节点分别指向x的子节点
            minNode.left = x.left;
            minNode.right = x.right;
            x = minNode;
        }else if(i < 0){
            x.left = remove(x.left, key);
        }else{
            x.right = remove(x.right, key);
        }

        return x;
    }

    public int size(){
        return count;
    }

    public boolean isEmpty(){
        return count == 0;
    }

    public K min(){
        BinaryNode<K, V> minNode = min(root);
        return minNode == null ? null : minNode.key;
    }
    private BinaryNode<K, V> min(BinaryNode<K, V> x){
        if(x == null){
            return null;
        }
        while(x.left != null){
            x = x.left;
        }
        return x;
    }

    public K max(){
        BinaryNode<K, V> maxNode = max(root);
        return maxNode == null ? null : maxNode.key;
    }
    private BinaryNode<K, V> max(BinaryNode<K, V> x){
        if(x == null){
            return null;
        }
        while(x.right != null){
            x = x.right;
        }
        return x;
    }

    /**
     * 前序遍历
     * @return
     */
    public Queue<K> preErgodic(){
        Queue<K> queue = new LinkedList<>();
        preErgodic(root, queue);
        return queue;
    }

    private void preErgodic(BinaryNode<K, V> x, Queue<K> queue) {
        if(x == null){
            return;
        }
        queue.add(x.key);
        preErgodic(x.left, queue);
        preErgodic(x.right, queue);
    }

    /**
     * 中序遍历
     * @return
     */
    public Queue<K> mindErgodic(){
        Queue<K> queue = new LinkedList<>();
        mindErgodic(root, queue);
        return queue;
    }

    private void mindErgodic(BinaryNode<K, V> x, Queue<K> queue) {
        if(x == null){
            return;
        }
        mindErgodic(x.left, queue);
        queue.add(x.key);
        mindErgodic(x.right, queue);
    }

    /**
     * 后序遍历
     * @return
     */
    public Queue<K> afterErgodic(){
        Queue<K> queue = new LinkedList<>();
        afterErgodic(root, queue);
        return queue;
    }

    private void afterErgodic(BinaryNode<K, V> x, Queue<K> queue) {
        if(x == null){
            return;
        }
        afterErgodic(x.left, queue);
        afterErgodic(x.right, queue);
        queue.add(x.key);
    }

    /**
     * 层序遍历
     * @return
     */
    public Queue<K> layerErgodic(){
        Queue<K> queue = new LinkedList<>();
        Queue<BinaryNode<K, V>> temp = new LinkedList<>();
        if(root == null){
            return queue;
        }
        temp.add(root);
        while(!temp.isEmpty()){
            BinaryNode<K, V> poll = temp.poll();
            queue.add(poll.key);
            if(poll.left != null){
                temp.add(poll.left);
            }
            if(poll.right != null){
                temp.add(poll.right);
            }
        }
        return queue;
    }

    /**
     * 二叉树最大深度
     * @return
     */
    public int maxDepth(){
        return maxDepth(root);
    }

    private int maxDepth(BinaryNode<K, V> x) {
        if(x == null){
            return 0;
        }
        int leftMaxDepth = maxDepth(x.left);
        int rightMaxDepth = maxDepth(x.right);
        return Math.max(leftMaxDepth, rightMaxDepth) + 1;
    }

    private static class BinaryNode<K, V>{
        final K key;
        V value;
        BinaryNode<K, V> left;
        BinaryNode<K, V> right;

        BinaryNode(K key, V value, BinaryNode<K, V> left, BinaryNode<K, V> right) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
        }
    }
}
